 /*
 *  RSA.cpp
 *  TubeTruck
 *
 *  Created by Aeturnum on 5/22/08.
 *  Copyright 2008 __MyCompanyName__. All rights reserved.
 *
 */

#include "RSA.h"

using namespace std;
using namespace TTP;

/*
struct RSA mKeys;
*/

void Crypt::RSA::makeHash() {
   vector<byte> temp(getPublicKey());
   
   mIdHash = *((uint32_t*)&temp[0]);
   for(uint32_t i = 4; i < temp.size(); i+=4) {
      mIdHash ^= *((uint32_t*)&temp[i]);
   }
}

void Crypt::RSA::generateKeys(int size, unsigned long e) {
   srand(clock());
   mKeys = RSA_generate_key(size, e, NULL, NULL);
   makeHash();
}

Crypt::RSA::RSA() {
   mKeys = RSA_new();
   vector<byte> e;
   
   ::RSA* tempKeys;
   srand(clock());
   tempKeys = RSA_generate_key(1024, 65537, NULL, NULL);
   
   e.resize(BN_num_bytes(tempKeys->e));
   BN_bn2bin(tempKeys->e, &e[0]);
   mKeys->e = BN_bin2bn(&e[0], e.size(),NULL);
   RSA_free(tempKeys);
}

Crypt::RSA::RSA(int size) {
   generateKeys(size, 65537);
}

Crypt::RSA::RSA(const RSA& o ) : mIdHash(o.mIdHash) {
   mKeys = o.mKeys;
   RSA_up_ref(mKeys);
}

Crypt::RSA::~RSA() {RSA_free(mKeys);}

vector<byte> Crypt::RSA::getPublicKey() {
   vector<byte> temp;
   temp.resize(BN_num_bytes(mKeys->n));
   BN_bn2bin(mKeys->n, &temp[0]);
   return temp;
}

vector<byte> Crypt::RSA::getPrivateKey() {
   vector<byte> temp;
   temp.resize(BN_num_bytes(mKeys->d));
   BN_bn2bin(mKeys->d, &temp[0]);
   return temp;   
}

uint32_t Crypt::RSA::getIdHash() {
   return mIdHash;
}

void Crypt::RSA::setPrivateKey(vector<byte>& d) {
   if(mKeys->d != NULL) {
      BN_free(mKeys->d);
      mKeys->d = 0;
   }
   mKeys->d = BN_bin2bn(&d[0], d.size(), NULL);   
   makeHash();
}

void Crypt::RSA::setPublicKey(vector<byte>& n) {
   if(mKeys->n != NULL) {
      BN_free(mKeys->n);
      mKeys->n = 0;
   }
   mKeys->n = BN_bin2bn(&n[0], n.size(), NULL);
}

vector<byte> Crypt::RSA::encrypt(const vector<byte>& data) {
   const uint32_t chunkSize = RSA_size(mKeys)-42;
   vector< vector<byte> > dataChunks;
   unsigned long length = data.size();
   vector<byte>::const_iterator itr = data.begin();
   
   do {
      vector<byte>::const_iterator end; 
      if(length > chunkSize) {
         end = itr + chunkSize;
         length -= chunkSize;
      } else {
         end = data.end();
         length = 0;
      }
      
      vector<byte> chunk(itr,end);
//      cout << "Encrypt chunk(" << chunk.size() << "):\n" << chunk;
//      cout << "Length: " << length << "\n"; 
      dataChunks.insert(dataChunks.end(), chunk);
      itr = itr + chunkSize;
   } while(length != 0);
   
   vector<byte> retVal;
   for(uint32_t i = 0; i< dataChunks.size(); i++) {
      vector<byte> encryptedChunk;
      int encryptedSize = -1;
      
      encryptedChunk.reserve(RSA_size(mKeys));
      encryptedChunk.resize(RSA_size(mKeys));
      
//      cout << "Encrypting:\n";
//      cout << "dataChunks[i](" << dataChunks[i].size() << "):\n" << dataChunks[i];
      
      encryptedSize = RSA_public_encrypt(dataChunks[i].size(), &dataChunks[i][0], &encryptedChunk[0], mKeys, RSA_PKCS1_OAEP_PADDING);
      
      if(encryptedSize == -1) {
         ERR_load_crypto_strings();
         char errorString[400];
         ERR_error_string(ERR_get_error(), errorString);
         printf("\nRSA::Encrytp: %s\n", errorString);
         return retVal;
      }
      
      encryptedChunk.resize(encryptedSize);
      
      retVal.insert(retVal.end(), encryptedChunk.begin(), encryptedChunk.end());
   }
   
//   cout << "retVal(" << retVal.size() << "):\n" << retVal;
   
   return retVal;
}
vector<byte> Crypt::RSA::decrypt(const vector<byte>& data) {
   
   const uint32_t chunkSize = RSA_size(mKeys);
   vector< vector<byte> > dataChunks;
   unsigned long length = data.size();
   vector<byte>::const_iterator itr = data.begin();
   
   do {
      vector<byte>::const_iterator end; 
      if(length > chunkSize) {
         end = itr + chunkSize;
         length -= chunkSize;
      } else {
         end = data.end();
         length = 0;
      }
      
      vector<byte> chunk(itr,end);
      dataChunks.insert(dataChunks.end(), chunk);
      itr = itr + chunkSize;
   } while(length != 0);
   
   vector<byte> retVal;
   for(uint32_t i = 0; i< dataChunks.size(); i++) {
      vector<byte> encryptedChunk;
      int plaintextSize = -1;
      
      encryptedChunk.reserve(RSA_size(mKeys));
      encryptedChunk.resize(RSA_size(mKeys));
      
      plaintextSize = RSA_private_decrypt(dataChunks[i].size(), &dataChunks[i][0], &encryptedChunk[0], mKeys, RSA_PKCS1_OAEP_PADDING);
      
      if(plaintextSize == -1) {
         ERR_load_crypto_strings();
         char errorString[400];
         ERR_error_string(ERR_get_error(), errorString);
         printf("\nRSA::decrypt: %s\n", errorString);
         return retVal;
      }
      
      encryptedChunk.resize(plaintextSize);
      
      retVal.insert(retVal.end(), encryptedChunk.begin(), encryptedChunk.end());
   }
   
   return retVal;
}

void Crypt::RSA::encrypt(Encryptable& data) {
   vector<byte> dataPlainText = data.getEncryptedSection();
   vector<byte> retVal = encrypt(dataPlainText);
   data.setEncryptedData(retVal);
}                         
   
void Crypt::RSA::decrypt(Encryptable& data) {
   vector<byte> dataPlainText = data.getEncryptedSection();
   vector<byte> retVal = decrypt(dataPlainText);
   data.setPlainText(retVal);
}

void Crypt::RSA::sign(Signable& p) {
   vector<byte> signiture(20), signedSigniture(512);
   uint32_t bytesSigned = 0;
   
   vector<byte> unsignedSignable(p.getUnsignedData());
   
   
   SHA1(&unsignedSignable[0], unsignedSignable.size(), &signiture[0]);
   
   if( !RSA_sign(NID_sha1, &signiture[0], 20, &signedSigniture[0], &bytesSigned, mKeys) ) {
      ERR_load_crypto_strings();
      char errorString[400];
      ERR_error_string(ERR_get_error(), errorString);
      printf("\nRSA::sign: %s\n", errorString);
      return;
   }
   
   Crypt::Signiture retVal(signedSigniture);
   retVal.resize(bytesSigned);
   
   assert(bytesSigned == 128); // if this isn't true, we've got problems
   
   p.setSigniture(retVal);
}

bool Crypt::RSA::verify(Signable& p) {
   Crypt::Signiture sig = p.getSigniture(), ourSigniture;
   
   assert(sig.size() == 128); // if this isn't true, we've got problems
   
   vector<byte> unsignedSignable(p.getUnsignedData());
   
   SHA1(&unsignedSignable[0], unsignedSignable.size(), &ourSigniture[0]);
   
   if( !RSA_verify(NID_sha1, &ourSigniture[0], 20, &sig[0], sig.size(), mKeys) ) {
      return false;
   }
   
   return true;
}